{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "HOA Network Selection Neurips Submission.ipynb",
      "provenance": [
        {
          "file_id": "1Tgl6WD5WM35CmitLzwfNItHKPnMOP0oF",
          "timestamp": 1622695496870
        },
        {
          "file_id": "1S4fmk97iAXQ_WrpWxIgPLbF1R8bx_MQU",
          "timestamp": 1618947957836
        }
      ],
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      }
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a8A8wlKjZRka"
      },
      "source": [
        "### HOA Network Selection for...\n",
        "##Efficiently Identifying Task Groupings for Multi-Task Learning\n",
        "\n",
        "Licensed under the Apache License, Version 2.0\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FC5hhwkPZUxh"
      },
      "source": [
        "### CelebA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RAUHKIa1N8-Y"
      },
      "source": [
        "# HOA Groupings on the Validation Set.\n",
        "revised_integrals = {'5_o_Clock_Shadow': {'5_o_Clock_Shadow': 92.28, 'Black_Hair': 91.73, 'Blond_Hair': 92.09, 'Brown_Hair': 92.16, 'Goatee': 92.34, \n",
        "\t\t\t\t\t   'Mustache': 92.30, 'No_Beard': 92.18, 'Rosy_Cheeks': 92.12, 'Wearing_Hat': 92.20}, \n",
        "\t\t\t\t\t 'Black_Hair': {'5_o_Clock_Shadow': 90.41, 'Black_Hair': 90.19, 'Blond_Hair': 90.29, 'Brown_Hair': 90.48, 'Goatee': 90.24, \n",
        "\t\t\t\t\t   'Mustache': 90.46, 'No_Beard': 90.30, 'Rosy_Cheeks': 90.22, 'Wearing_Hat': 90.34}, \n",
        "\t\t\t\t\t 'Blond_Hair': {'5_o_Clock_Shadow': 94.87, 'Black_Hair': 95.00, 'Blond_Hair': 94.96, 'Brown_Hair': 94.87, 'Goatee': 94.95, \n",
        "\t\t\t\t\t   'Mustache': 94.94, 'No_Beard': 94.94, 'Rosy_Cheeks': 94.75, 'Wearing_Hat': 94.98}, \n",
        "\t\t\t\t\t 'Brown_Hair': {'5_o_Clock_Shadow': 84.44, 'Black_Hair': 84.42, 'Blond_Hair': 84.39, 'Brown_Hair': 84.39, 'Goatee': 84.09, \n",
        "\t\t\t\t\t   'Mustache': 84.31, 'No_Beard': 84.30, 'Rosy_Cheeks': 84.39, 'Wearing_Hat': 84.39}, \n",
        "\t\t\t\t\t 'Goatee': {'5_o_Clock_Shadow': 95.84, 'Black_Hair': 95.76, 'Blond_Hair': 95.81, 'Brown_Hair': 95.77, 'Goatee': 95.90, \n",
        "\t\t\t\t\t   'Mustache': 95.87, 'No_Beard': 96.07, 'Rosy_Cheeks': 95.83, 'Wearing_Hat': 95.83}, \n",
        "\t\t\t\t\t 'Mustache': {'5_o_Clock_Shadow': 95.69, 'Black_Hair': 95.68, 'Blond_Hair': 95.80, 'Brown_Hair': 95.77, 'Goatee': 95.75, \n",
        "\t\t\t\t\t   'Mustache': 95.84, 'No_Beard': 95.75, 'Rosy_Cheeks': 95.82, 'Wearing_Hat': 95.74}, \n",
        "\t\t\t\t\t 'No_Beard': {'5_o_Clock_Shadow': 94.14, 'Black_Hair': 94.01, 'Blond_Hair': 93.99, 'Brown_Hair': 94.17, 'Goatee': 94.47, \n",
        "\t\t\t\t\t   'Mustache': 94.30, 'No_Beard': 94.44, 'Rosy_Cheeks': 93.96, 'Wearing_Hat': 94.14}, \n",
        "\t\t\t\t\t 'Rosy_Cheeks': {'5_o_Clock_Shadow': 94.66, 'Black_Hair': 94.69, 'Blond_Hair': 94.64, 'Brown_Hair': 94.72, 'Goatee': 94.68, \n",
        "\t\t\t\t\t   'Mustache': 94.62, 'No_Beard': 94.62, 'Rosy_Cheeks': 94.64, 'Wearing_Hat': 94.59}, \n",
        "\t\t\t\t\t 'Wearing_Hat': {'5_o_Clock_Shadow': 98.61, 'Black_Hair': 98.63, 'Blond_Hair': 98.64, 'Brown_Hair': 98.56, 'Goatee': 98.69, \n",
        "\t\t\t\t\t   'Mustache': 98.66, 'No_Beard': 98.67, 'Rosy_Cheeks': 98.59, 'Wearing_Hat': 98.69}}\n",
        "\n",
        "# Minimize Error rather than Maximize Accuracy\n",
        "for key in revised_integrals:\n",
        "\tfor task in revised_integrals[key]:\n",
        "\t\trevised_integrals[key][task] = 100.0 - revised_integrals[key][task]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EBmh3dFZZem0"
      },
      "source": [
        "### Taskonomy"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Rh3BYd--heV"
      },
      "source": [
        "# HOA Groupings on the Validation Set.\n",
        "revised_integrals = {'s': {'s': 0.0509, 'd': 0.0528, 'n': 0.0471, 't': 0.0606, 'k': 0.0655},\n",
        "                     'd': {'s': 0.2636, 'd': 0.2616, 'n': 0.2551, 't': 0.2675, 'k': 0.2674},\n",
        "                     'n': {'s': 0.1004, 'd': 0.1002, 'n': 0.0975, 't': 0.1082, 'k': 0.1110},\n",
        "                     't': {'s': 0.0287, 'd': 0.0423, 'n': 0.0310, 't': 0.0337, 'k': 0.0232},\n",
        "                     'k': {'s': 0.0972, 'd': 0.1068, 'n': 0.0568, 't': 0.0933, 'k': 0.0910}}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dlChfA6SZi01"
      },
      "source": [
        "### Network Selection Algorithm"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DhdVOK1lNqzC"
      },
      "source": [
        "import math \n",
        "import numpy as np\n",
        "\n",
        "def gen_task_combinations(tasks, rtn, index, path, path_dict):\n",
        "  if index >= len(tasks):\n",
        "    return \n",
        "\n",
        "  for i in range(index, len(tasks)):\n",
        "    cur_task = tasks[i]\n",
        "    new_path = path\n",
        "    new_dict = {k:v for k,v in path_dict.items()}\n",
        "    \n",
        "    # Building from a tree with two or more tasks...\n",
        "    if new_path:\n",
        "      new_dict[cur_task] = 0.\n",
        "      for prev_task in path_dict:\n",
        "        new_dict[prev_task] += revised_integrals[prev_task][cur_task]\n",
        "        new_dict[cur_task] += revised_integrals[cur_task][prev_task]\n",
        "      new_path = '{}|{}'.format(new_path, cur_task)\n",
        "      rtn[new_path] = new_dict\n",
        "    else: # First element in a new-formed tree\n",
        "      new_dict[cur_task] = 0.\n",
        "      new_path = cur_task\n",
        "\n",
        "    gen_task_combinations(tasks, rtn, i+1, new_path, new_dict)\n",
        "\n",
        "    # Fix single-task accuracy since dfs is finished for this task.\n",
        "    if '|' not in new_path:\n",
        "      new_dict[cur_task] = revised_integrals[cur_task][cur_task]\n",
        "      rtn[new_path] = new_dict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NuRJiEmxWKMZ"
      },
      "source": [
        "rtn = {}\n",
        "tasks = list(revised_integrals.keys())\n",
        "num_tasks = len(tasks)\n",
        "task_combinations = gen_task_combinations(tasks=tasks, rtn=rtn, index=0, path='', path_dict={})\n",
        "\n",
        "# Normalize by the number of times the accuracy of any given element has been summed. \n",
        "# i.e. (a,b,c) => [acc(a|b) + acc(a|c)]/2 + [acc(b|a) + acc(b|c)]/2 + [acc(c|a) + acc(c|b)]/2\n",
        "for group in rtn:\n",
        "  if '|' in group:\n",
        "    for task in rtn[group]:\n",
        "      rtn[group][task] /= (len(group.split('|')) - 1)\n",
        "\n",
        "print(rtn)\n",
        "assert(len(rtn.keys()) == 2**len(revised_integrals.keys()) - 1)\n",
        "rtn_tup = [(key,val) for key,val in rtn.items()]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WvTgeaR3KWMd"
      },
      "source": [
        "def select_groups(index, cur_group, best_group, best_val, splits):\n",
        "  # Check if this group covers all tasks.\n",
        "  task_set = set()\n",
        "  for group in cur_group:\n",
        "    for task in group.split('|'): task_set.add(task)\n",
        "  if len(task_set) == num_tasks:\n",
        "    best_tasks = {task:1e6 for task in task_set}\n",
        "    \n",
        "    # Compute the per-task best scores for each task and average them together.\n",
        "    for group in cur_group:\n",
        "      for task in cur_group[group]:\n",
        "        # Minimize error.\n",
        "        best_tasks[task] = min(best_tasks[task], cur_group[group][task])\n",
        "    group_avg = np.mean(list(best_tasks.values()))\n",
        "    \n",
        "    # Compare with the best grouping seen thus far.\n",
        "    if group_avg < best_val[0]:\n",
        "      print(cur_group)\n",
        "      best_val[0] = group_avg\n",
        "      best_group.clear()\n",
        "      for entry in cur_group:\n",
        "        best_group[entry] = cur_group[entry]\n",
        "  \n",
        "  # Base case.\n",
        "  if len(cur_group.keys()) == splits:\n",
        "    return\n",
        "\n",
        "  # Back to combinatorics \n",
        "  for i in range(index, len(rtn_tup)):\n",
        "    selected_group, selected_dict = rtn_tup[i]\n",
        "\n",
        "    new_group = {k:v for k,v in cur_group.items()}\n",
        "    new_group[selected_group] = selected_dict\n",
        "\n",
        "    if len(new_group.keys()) <= splits:\n",
        "      select_groups(i + 1, new_group, best_group, best_val, splits)\n",
        "\n",
        "selected_group = {}\n",
        "selected_val = [100000000]\n",
        "select_groups(index=0, cur_group={}, best_group=selected_group, best_val=selected_val, splits=2)\n",
        "print(selected_group)\n",
        "print(selected_val)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}